﻿using System;
using System.Collections.Generic;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;

namespace Mars.GpuInstance
{
    public static class InstancerUtil
    {
        
        public static readonly string SHADER_UNITY_INTERNAL_ERROR = "Hidden/InternalErrorShader";
        private static readonly string DEFAULT_SHADER_PATH = "Assets/res/shader/gpu_instancer/";

        public static bool IsShaderInstanced(Shader shader)
        {
#if UNITY_EDITOR
            if (shader == null || shader.name == SHADER_UNITY_INTERNAL_ERROR)
            {
                return false;
            }
            string originalAssetPath = AssetDatabase.GetAssetPath(shader);
            string originalShaderText = "";
            try
            {
                originalShaderText = System.IO.File.ReadAllText(originalAssetPath);
            }
            catch (Exception)
            {
                return false;
            }
            if (!string.IsNullOrEmpty(originalShaderText))
            {
                if (originalAssetPath.ToLower().EndsWith(".shadergraph"))
                {
                    return false;
                }
                return originalShaderText.Contains("GPUInstancing_indirect.cginc");
            }
#endif
            return false;
        }
        
        public static Shader CreateInstancedShader(Shader originalShader, bool useOriginal = false)
        {
#if UNITY_EDITOR
            try
            {
                if (originalShader == null || originalShader.name == SHADER_UNITY_INTERNAL_ERROR)
                {
                    return null;
                }
                
                Shader originalShaderRef = Shader.Find(originalShader.name);
                string originalAssetPath = AssetDatabase.GetAssetPath(originalShaderRef);
                string extension = System.IO.Path.GetExtension(originalAssetPath);

                // can not work with ShaderGraph or other non shader code
                if (extension != ".shader")
                {
                    if (extension.EndsWith("pack"))
                    {
                        return CreateInstancedShaderPack(originalShader, originalAssetPath);
                    }
                    else
                    {
                        return null;
                    }
                }

                EditorUtility.DisplayProgressBar("生成支持GPU Instancer的shader", "生成 " + originalShader.name + "...", 0.1f);

                #region Remove Existing procedural setup
                System.Text.StringBuilder sb = new System.Text.StringBuilder();
                using (System.IO.StreamReader sr = new System.IO.StreamReader(originalAssetPath))
                {
                    while (!sr.EndOfStream)
                    {
                        string line = sr.ReadLine();
                        if (!line.Contains("#pragma instancing_options") && !line.Contains("GPUInstancing_indirect.cginc") && !line.Contains("#pragma multi_compile_instancing"))
                        {
                            sb.Append(line + "\n");
                        }
                    }
                }
                string originalShaderText = sb.ToString();
                #endregion Remove Existing procedural setup

                // create shader versions for packages inside GPUI folder
                bool createInDefaultFolder = originalAssetPath.StartsWith("Packages/");

                // Packages/com.unity.render-pipelines.universal/Shaders/Lit.shader
                // if URP, replace relative paths
                bool isURP = false;
                string urpIncludeAddition = "Packages/com.unity.render-pipelines.universal/";
                if (originalShader.name.StartsWith("Universal Render Pipeline/"))
                {
                    isURP = true;
                    string[] urpSplit = originalAssetPath.Split('/');
                    bool foundURP = false;
                    for (int i = 0; i < urpSplit.Length; i++)
                    {
                        if (urpSplit[i].Contains(".shader"))
                        {
                            break;
                        }
                        if (foundURP)
                        {
                            urpIncludeAddition += urpSplit[i] + "/";
                        }
                        else
                        {
                            if (urpSplit[i] == "com.unity.render-pipelines.universal")
                            {
                                foundURP = true;
                            }
                        }
                    }
                }

                EditorUtility.DisplayProgressBar("生成支持GPU Instancer的shader", "生成 " + originalShader.name + "...", 0.5f);

                string newShaderName = (useOriginal && !createInDefaultFolder) ? "" : "GPUInstancer/" + originalShader.name;
                string newShaderText = (useOriginal && !createInDefaultFolder) ? originalShaderText.Replace("\r\n", "\n") : originalShaderText.Replace("\r\n", "\n").Replace(originalShader.name, newShaderName);

                string includePath = GetShaderIncludePath(originalAssetPath, createInDefaultFolder, out originalAssetPath);

                int lastIndex = 0;
                string searchStart = "";
                string searchEnd = "";
                int foundIndex = -1;

                // For vertex/fragment and surface shaders
                newShaderText = ReplaceShaderTextCGProgram(includePath, newShaderText, false);
                
                // For URP Shaders Include relative path fix
                #region URP relative path fix
                if (isURP && createInDefaultFolder)
                {
                    lastIndex = 0;
                    searchStart = "#include \"";
                    searchEnd = "\"";
                    string restOfText;

                    foundIndex = -1;
                    while (true)
                    {
                        foundIndex = newShaderText.IndexOf(searchStart, lastIndex);
                        if (foundIndex == -1)
                            break;
                        lastIndex = foundIndex + searchStart.Length + 1;

                        restOfText = newShaderText.Substring(foundIndex + searchStart.Length, newShaderText.Length - foundIndex - searchStart.Length);
                        if (!restOfText.StartsWith("URP") && !restOfText.StartsWith("CoreRP") && !restOfText.StartsWith("Packages"))
                        {
                            newShaderText = newShaderText.Substring(0, foundIndex + searchStart.Length) + urpIncludeAddition + restOfText;
                            lastIndex += urpIncludeAddition.Length;
                        }

                        foundIndex = newShaderText.IndexOf(searchEnd, lastIndex);
                        lastIndex = foundIndex;
                    }
                }
                #endregion URP relative path fix

                // For SRP Shaders
                newShaderText = ReplaceShaderTextHLSLProgram(includePath, newShaderText, false);

                Shader instancedShader = SaveInstancedShader(originalAssetPath, useOriginal, extension, newShaderText, originalShader, newShaderName);

                if (instancedShader != null)
                {
                    Debug.LogError("创建支持gpu instancer版本的shader成功: " + originalShader.name, instancedShader);
                }
                EditorUtility.ClearProgressBar();

                return instancedShader;
            }
            catch (Exception e)
            {
                if (e is System.IO.DirectoryNotFoundException && e.Message.ToLower().Contains("unity_builtin_extra"))
                {
                    Debug.LogError("\"" + originalShader.name + "\" 是 shader built-in shader, 没法自动生成"); 
                }
                else
                {
                    Debug.LogError(e);
                }
                EditorUtility.ClearProgressBar();
            }
#endif
            return null;
        }
        
        public static Shader CreateInstancedShaderPack(Shader originalShader, string originalAssetPath)
        {
#if UNITY_EDITOR
            EditorUtility.DisplayProgressBar("生成支持GPU Instancer的shader", "生成 " + originalShader.name + "...", 0.1f);
            string extension = System.IO.Path.GetExtension(originalAssetPath);
            #region Remove Existing procedural setup
            System.Text.StringBuilder sb = new System.Text.StringBuilder();
            using (System.IO.StreamReader sr = new System.IO.StreamReader(originalAssetPath))
            {
                while (!sr.EndOfStream)
                {
                    string line = sr.ReadLine();
                    string[] lines = line.Split(new string[] { "\\n" }, StringSplitOptions.None);
                    for (int i = 0; i < lines.Length; i++)
                    {
                        line = lines[i];
                        if (!line.Contains("#pragma instancing_options") && !line.Contains("GPUInstancing_indirect.cginc") && !line.Contains("#pragma multi_compile_instancing"))
                        {
                            if (i == lines.Length - 1)
                            {
                                sb.Append(line);
                            }
                            else
                            {
                                sb.Append(line + "\\n");
                            }
                        }
                    }
                }
            }
            string originalShaderText = sb.ToString();
            #endregion Remove Existing procedural setup
            // create shader versions for packages inside GPUI folder
            bool createInDefaultFolder = originalAssetPath.StartsWith("Packages/");

            EditorUtility.DisplayProgressBar("生成支持GPU Instancer的shader", "生成 " + originalShader.name + "...", 0.5f);

            string newShaderName = "GPUInstancer/" + originalShader.name;
            string newShaderText = originalShaderText.Replace("\r\n", "\n").Replace(originalShader.name, newShaderName);

            string includePath = GetShaderIncludePath(originalAssetPath, createInDefaultFolder, out originalAssetPath);

            // For vertex/fragment and surface shaders
            newShaderText = ReplaceShaderTextCGProgram(includePath, newShaderText, true);

            // For SRP Shaders
            newShaderText = ReplaceShaderTextHLSLProgram(includePath, newShaderText, true);

            Shader instancedShader = SaveInstancedShader(originalAssetPath, false, extension, newShaderText, originalShader, newShaderName);

            if (instancedShader != null)
            {
                Debug.LogError("生成 Gpu instance版本shader失败，原shader: " + originalShader.name, instancedShader);
            }
            EditorUtility.ClearProgressBar();

            return instancedShader;
#else
            return null;
#endif
        }
        
        private static string GetShaderIncludePath(string originalAssetPath, bool createInDefaultFolder, out string newAssetPath)
        {
            newAssetPath = originalAssetPath;
            string[] oapSplit = originalAssetPath.Split('/');
            if (createInDefaultFolder)
            {
                if (!System.IO.Directory.Exists(DEFAULT_SHADER_PATH))
                {
                    System.IO.Directory.CreateDirectory(DEFAULT_SHADER_PATH);
                }
                newAssetPath = DEFAULT_SHADER_PATH + "GPUI_"+ oapSplit[oapSplit.Length - 1];
            }
            return "Packages/com.txcombo.c1.simple-gpu-instance/Shader/Include/GPUInstancing_indirect.cginc";
        }
        
        private static string ReplaceShaderTextCGProgram(string includePath, string newShaderText, bool doubleEscape)
        {
            // For vertex/fragment and surface shaders
            #region CGPROGRAM
            int lastIndex = 0;
            string searchStart = "CGPROGRAM";
            string additionTextStart = "\n#include \"UnityCG.cginc\"\n#include \"" + includePath + "\"\n#pragma instancing_options procedural:setup\n#pragma multi_compile_instancing";
            if (doubleEscape)
            {
                additionTextStart = additionTextStart.Replace("\n", "\\n").Replace("\"", "\\\"");
            }
            string searchEnd = "ENDCG";
            string additionTextEnd = "";//"#include \"" + includePath + "\"\n";

            int foundIndex = -1;
            while (true)
            {
                foundIndex = newShaderText.IndexOf(searchStart, lastIndex);
                if (foundIndex == -1)
                    break;
                lastIndex = foundIndex + searchStart.Length + additionTextStart.Length + 1;

                newShaderText = newShaderText.Substring(0, foundIndex + searchStart.Length) + additionTextStart + newShaderText.Substring(foundIndex + searchStart.Length, newShaderText.Length - foundIndex - searchStart.Length);

                foundIndex = newShaderText.IndexOf(searchEnd, lastIndex);
                lastIndex = foundIndex + searchStart.Length + additionTextEnd.Length + 1;
                newShaderText = newShaderText.Substring(0, foundIndex) + additionTextEnd + newShaderText.Substring(foundIndex, newShaderText.Length - foundIndex);
            }
            #endregion CGPROGRAM

            return newShaderText;
        }
        
        private static string ReplaceShaderTextHLSLProgram(string includePath, string newShaderText, bool doubleEscape)
        {
            // For SRP Shaders
            #region HLSLPROGRAM
            int lastIndex = 0;
            string searchStart = "HLSLPROGRAM";
            string additionTextStart = "\n#include \"" + includePath + "\"\n#pragma instancing_options procedural:setup\n#pragma multi_compile_instancing\n";
            additionTextStart = "\n#include \"Packages/com.unity.render-pipelines.universal/ShaderLibrary/Core.hlsl\"\n" + additionTextStart;
            if (doubleEscape)
            {
                additionTextStart = additionTextStart.Replace("\n", "\\n").Replace("\"", "\\\"");
            }
            string searchEnd = "ENDHLSL";
            string additionTextEnd = "";
            if (doubleEscape)
            {
                additionTextEnd = additionTextEnd.Replace("\n", "\\n").Replace("\"", "\\\"");
            }
            int foundIndex = -1;
            while (true)
            {
                foundIndex = newShaderText.IndexOf(searchStart, lastIndex);
                if (foundIndex == -1)
                    break;
                lastIndex = foundIndex + searchStart.Length + additionTextStart.Length + 1;

                newShaderText = newShaderText.Substring(0, foundIndex + searchStart.Length) + additionTextStart + newShaderText.Substring(foundIndex + searchStart.Length, newShaderText.Length - foundIndex - searchStart.Length);

                foundIndex = newShaderText.IndexOf(searchEnd, lastIndex);
                lastIndex = foundIndex + searchStart.Length + additionTextEnd.Length + 1;
                newShaderText = newShaderText.Substring(0, foundIndex) + additionTextEnd + newShaderText.Substring(foundIndex, newShaderText.Length - foundIndex);
            }
            #endregion HLSLPROGRAM

            #region Temp Unity 2021.2.8-11 Bug Fix
#if UNITY_2021_2_8 || UNITY_2021_2_9 || UNITY_2021_2_10 || UNITY_2021_2_11
            if (!GPUInstancerConstants.gpuiSettings.IsStandardRenderPipeline())
                newShaderText = newShaderText.Replace("HLSLPROGRAM", "HLSLPROGRAM\n#pragma multi_compile _ PROCEDURAL_INSTANCING_ON").Replace("#pragma multi_compile_instancing", "");
#endif
            #endregion Temp Unity 2021.2.8 Bug Fix

            return newShaderText;
        }
        
        private static Shader SaveInstancedShader(string originalAssetPath, bool useOriginal, string extension, string newShaderText, Shader originalShader, string newShaderName)
        {
#if UNITY_EDITOR

            string originalFileName = System.IO.Path.GetFileName(originalAssetPath);
            string newAssetPath = useOriginal ? originalAssetPath : originalAssetPath.Replace(originalFileName, originalFileName.Replace(extension, "_GPUI" + extension));
            byte[] bytes = System.Text.Encoding.UTF8.GetBytes(newShaderText);
            System.IO.FileStream fs = System.IO.File.Create(newAssetPath);
            fs.Write(bytes, 0, bytes.Length);
            fs.Close();
            EditorUtility.DisplayProgressBar("生成支持GPU Instancer的shader", "生成 " + originalShader.name + "...", 0.8f);
            AssetDatabase.ImportAsset(newAssetPath, ImportAssetOptions.ForceUpdate);
            AssetDatabase.Refresh();

            Shader instancedShader = AssetDatabase.LoadAssetAtPath<Shader>(newAssetPath);
            if (instancedShader == null)
            {
                instancedShader = Shader.Find(newShaderName);
            }

            return instancedShader;
#else
            return null;
#endif
        }
        
        
        public static Mesh GenerateMeshWithUVs(Mesh originalMesh, GPUSkinnedMeshData smd)
        {
            List<Vector4> boneIndexes = new List<Vector4>();
            List<Vector4> boneWeights = new List<Vector4>();
            // xy: most important bone index, zw: non proportional scale bone index
            Vector4 boneIndexVector = Vector4.zero;
            Vector4 boneWeightVector = Vector4.zero;
            foreach (BoneWeight boneWeight in originalMesh.boneWeights)
            {
                boneIndexVector.x = smd.boneIndexes[boneWeight.boneIndex0];
                boneIndexVector.y = smd.boneIndexes[boneWeight.boneIndex1];
                
                //boneIndexVector.z = nonPropScaleBoneIndexes.FindIndex((bi)=> Mathf.Approximately(bi,boneIndexVector.x));
                //boneIndexVector.w = nonPropScaleBoneIndexes.FindIndex((bi)=> Mathf.Approximately(bi,boneIndexVector.y));
                boneIndexes.Add(boneIndexVector);

                boneWeightVector.x = boneWeight.weight0;
                boneWeightVector.y = boneWeight.weight1;
                boneWeightVector.z = boneWeight.weight2;
                boneWeightVector.w = boneWeight.weight3;
                boneWeights.Add(boneWeightVector);
            }
            Mesh mesh = new Mesh();
            mesh.name = originalMesh.name + "_gpu";
            mesh.subMeshCount = originalMesh.subMeshCount;
            mesh.vertices = originalMesh.vertices;

            for (int i = 0; i < mesh.subMeshCount; i++)
            {
                mesh.SetTriangles(originalMesh.GetTriangles(i), i);
            }

            mesh.normals = originalMesh.normals;
            mesh.tangents = originalMesh.tangents;
            mesh.colors = originalMesh.colors;
            mesh.uv = originalMesh.uv;
            mesh.uv2 = originalMesh.uv2;
            mesh.SetUVs(2, boneIndexes);
            mesh.SetUVs(3, boneWeights);
            mesh.UploadMeshData(true);
            
            ApplyMeshStreamLayout(mesh);
            
            return mesh;
        }
        
        public static int GetStringHashCode(string value)
        {
            return Animator.StringToHash(value);
        }

        public static void ApplyMeshStreamLayout(Mesh mesh)
        {
            // mesh stream layout
            // steam 0: position, uv2, uv3, 计算屏幕空间顶点坐标需要的数据
            List<VertexAttributeDescriptor> layoutList = new List<VertexAttributeDescriptor>();
            layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.Position,
                mesh.GetVertexAttributeFormat(VertexAttribute.Position),
                mesh.GetVertexAttributeDimension(VertexAttribute.Position),
                stream: 0));
            if (mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord2) > 0)
            {
                layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.TexCoord2,
                    mesh.GetVertexAttributeFormat(VertexAttribute.TexCoord2),
                    mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord2),
                    stream: 0));
            }
            if (mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord3) > 0)
            {
                layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.TexCoord3,
                    mesh.GetVertexAttributeFormat(VertexAttribute.TexCoord3),
                    mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord3),
                    stream: 0));
            }
            
            // stream 1: normal, tangent, color, uv0, uv1, 与屏幕空间顶点坐标计算无关的数据
            layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.Normal,
                mesh.GetVertexAttributeFormat(VertexAttribute.Normal),
                mesh.GetVertexAttributeDimension(VertexAttribute.Normal),
                stream: 1));
            layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.Tangent,
                mesh.GetVertexAttributeFormat(VertexAttribute.Tangent),
                mesh.GetVertexAttributeDimension(VertexAttribute.Tangent),
                stream: 1));
            if (mesh.GetVertexAttributeDimension(VertexAttribute.Color) > 0)
            {
                layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.Color,
                    mesh.GetVertexAttributeFormat(VertexAttribute.Color),
                    mesh.GetVertexAttributeDimension(VertexAttribute.Color),
                    stream: 1));
            }
            if (mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord0) > 0)
            {
                layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.TexCoord0,
                    mesh.GetVertexAttributeFormat(VertexAttribute.TexCoord0),
                    mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord0),
                    stream: 1));
            }
            if (mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord1) > 0)
            {
                layoutList.Add(new VertexAttributeDescriptor(VertexAttribute.TexCoord1,
                    mesh.GetVertexAttributeFormat(VertexAttribute.TexCoord1),
                    mesh.GetVertexAttributeDimension(VertexAttribute.TexCoord1),
                    stream: 1));
            }
            
            mesh.SetVertexBufferParams(mesh.vertexCount, layoutList.ToArray());
        }

#if UNITY_EDITOR
        [MenuItem("GPUAnimatorTools/ApplyMeshStreamLayout")]
        public static void FixAllGPUAnimMeshStream()
        {
            var guids = AssetDatabase.FindAssets("_gpu t:mesh", new[] { "Assets/res" });
            foreach (var guid in guids)
            {
                string path = AssetDatabase.GUIDToAssetPath(guid);
                if(!path.ToLower().Contains("/gpuskin/")) continue; 
                Mesh mesh = AssetDatabase.LoadAssetAtPath<Mesh>(path);
                ApplyMeshStreamLayout(mesh);
                EditorUtility.SetDirty(mesh);
                AssetDatabase.ImportAsset(path);
            }
        }
#endif
    }
}
